(* Written by Zachary Salzbank *)

open Ast
open Sast

module NameMap = Map.Make(struct
  type t = string
  let compare x y = Pervasives.compare x y
end)

type node = {
  value: int;
  child: int array;
}

type environment = {
  locals: int NameMap.t;
  globals: int NameMap.t;
  nodes: node array;
}

exception ReturnException of int * environment 

let node_check i =
  if i < 0 then
    raise (Failure("node does not exist yet"))
  else
    i

let get_child a i = 
  let cur_len = Array.length a
  in if i < cur_len then
    a, Array.get a i
  else
    let a = Array.append a (Array.make (i - cur_len + 1)  (-1))
    in a, Array.get a i

let array_replace a i x = 
  let left = Array.sub a 0 i
  in let right = Array.sub a (i+1) ((Array.length a)-i-1)
  in let nodes = Array.append left (Array.make 1 x)
  in Array.append nodes right

(* Main entry point: run a program *)

let run (vars, funcs) =
  (* Put function declarations in a symbol table *)
  let func_decls = List.fold_left
    (fun funcs fdecl -> NameMap.add fdecl.ffname fdecl funcs)
    NameMap.empty
    funcs

  in let rec l_value (env:environment) = function
        Id(var) ->
          if NameMap.mem var.vvname env.locals then
            (NameMap.find var.vvname env.locals), env
          else
            (NameMap.find var.vvname env.globals), env
      | Unop(lv, op) ->
          let lvi, env = l_value env (fst lv)
          in let rnode = Array.get env.nodes (node_check lvi)
          in match op with
              Child(e) ->
                let v, env = eval env e
                in let children, child = get_child rnode.child (node_check v)
                in let rnode' = {rnode with child = children}
                in let nodes = array_replace env.nodes lvi rnode'
                in let env' = {env with nodes = nodes}
                in child, env'
            | ValueOf  ->
                rnode.value, env
  and eval (env:environment) e = 
   (* Evaluate an expression and return (value, updated environment) *)
      let e, t = e in match e with
        Literal(c) -> (match c with 
            Integer(i) -> i
          | Character(ch) -> Char.code ch
          | Boolean(b) -> if b then 1 else 0
          | Null -> -1 
          ), env
      | Noexpr -> 1, env (* must be non-zero for the for loop predicate *)
      | Binop(e1, op, e2) ->
          let v1, env = eval env e1 in
          let v2, env = eval env e2 in
          let b1 = if (v1 == 1) then true else false in
          let b2 = if (v2 == 1) then true else false in
          let boolean i = if i then 1 else 0 in
          (match op with
              Add -> v1 + v2
            | Sub -> v1 - v2
            | Mult -> v1 * v2
            | Div -> v1 / v2
            | Equal -> boolean (v1 = v2)
            | Neq -> boolean (v1 != v2)
            | Less -> boolean (v1 < v2)
            | Leq -> boolean (v1 <= v2)
            | Greater -> boolean (v1 > v2)
            | Geq -> boolean (v1 >= v2)
            | BoolAnd -> boolean (b1 && b2)
            | BoolOr -> boolean (b1 || b2)
          ), env
      | Assign(lval, e) ->
          let v, env = eval env e
          in (match fst lval with
              Id(var) -> 
                if NameMap.mem var.vvname env.locals then(
                  let locals = NameMap.add var.vvname v env.locals in
                  v, {env with locals = locals;}
                ) else
                  let globals = NameMap.add var.vvname v env.globals
                  in v, {env with globals = globals;}
            | Unop(lv, op) ->
                let lvi, env = l_value env (fst lv)
                in let rnode = Array.get env.nodes (node_check lvi)
                in let n, env = (match op with
                    Child(e) ->
                      let iv, env = eval env e
                      in let a, cv = get_child rnode.child iv
                      in let a = array_replace a iv v
                      in {rnode with child = a}, env 
                  | ValueOf  ->
                      {rnode with value = v}, env)
                in let nodes = array_replace env.nodes lvi n
                in let env = {env with nodes = nodes;}
                in v, env)
      | Neg(e) -> 
          let v, env = eval env e
          in (v * -1), env
      | Bang(e) ->
          let v, env = eval env e
          in let value = if v = 1 then 0 else 1
          in value, env
      | Node(e) ->
          let nvalue, env = eval env (e, NullType) 
          in let n = {
            value = nvalue;
            child = Array.make 10 (-1);
          }
          in let index = Array.length env.nodes in
          let nodes = Array.append env.nodes (Array.make 1 n)
          in let env = {env with nodes = nodes;}
          in index, env
      | LValue(l) -> l_value env (fst l) 
      | Call(f, actuals) ->
          if f.ffname = "print" then
            let fml = List.hd f.fformals
            in let t = fml.vvtype
            in let v, env = eval env (List.hd actuals)
            in let s = match t with
                IntType -> string_of_int v
              | BooleanType -> if v==1 then "true" else "false"
              | CharType -> string_of_char (Char.chr v)
              | _ -> raise (Failure("Invalid print call"))
            in print_endline s;
            0, env
          else
            let fdecl = NameMap.find f.ffname func_decls
            in 
            let actuals, env = List.fold_left
              (fun (actuals, env) actual ->
                let v, env = eval env actual in
                v :: actuals, env)
              ([], env)
              (List.rev actuals)
          in try
            let env' = call fdecl actuals env
            in 0, {env' with locals = env.locals;}
          with ReturnException(v, env') ->
            let env = {env' with locals = env.locals;}
            in v, env 
  and add_local env vdecl =
    let dv = match vdecl.vvtype with
        NodeType(_) -> -1
      | _ -> 0
    in let value, env = match vdecl.vvdefault with
        Some(x) -> eval env x
      | None -> dv, env
    in 
    let locals = NameMap.add vdecl.vvname value env.locals
    in let env = {env with locals = locals;}
    in
    env
  and exec env = function
      Block(vars, stmts) -> 
        let env' = List.fold_left add_local env vars
        in let env' = List.fold_left exec env' stmts 
        in let locals = NameMap.fold 
          (fun lname lval lst -> if NameMap.mem lname env.locals then
              NameMap.add lname lval lst
            else
              lst
          )
          env'.locals
          NameMap.empty
        in {env' with locals = locals};
    | Expr(e) -> 
        let _, env = eval env e in env
    | If(e, s1, s2) ->
      let v, env = eval env e in
      exec env (if v != 0 then s1 else s2)
    | While(e, s) ->
      let rec loop env =
        let v, env = eval env e in
          if v != 0 then loop (exec env s) else env
      in loop env
    | Return(e) ->
      let v, env = eval env e in
      raise (ReturnException(v, env))

  (* Invoke a function and return an updated environment *)
  and call fdecl actuals (env : environment) =
    (* Enter the function: bind actual values to formal arguments *)
    let locals =
      List.fold_left2
          (fun locals formal actual -> NameMap.add formal actual locals)
          NameMap.empty
          (List.map (fun v -> v.vvname) fdecl.fformals)
          actuals
    in let env = {env with locals = locals}
    (* Initialize local variables *)
    in let env = List.fold_left
      add_local
      env 
      (List.rev fdecl.flocals)
    in
    (* Execute each statement in sequence, return updated environment *)
    List.fold_left exec env fdecl.fbody
 
  (* Run a program: initialize environment with globals, find and run "root" *)
  in let env = {
    globals = NameMap.empty;
    locals = NameMap.empty;
    nodes = Array.make 0 {
      value = 0;
      child = Array.make 10 (-1);
    };
  }
  in let env = List.fold_left
      (fun env vdecl ->
        let value, env = match vdecl.vvdefault with
            Some(x) -> eval env x
          | None -> 0, env
        in let globals = NameMap.add vdecl.vvname value env.globals
        in {env with globals = globals;}
      )
      env vars
  in try
    call (NameMap.find "root" func_decls) [] env
  with ReturnException(v, env) ->
    ignore(print_endline("Exited with code " ^ string_of_int v));
    env
